# encoding: utf-8
#!/usr/bin/python3
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.nn.modules.loss import _Loss




def multi_iou(pred,label,num_class=10):
    ## B*W*H
    pred = pred + 1
    label = label + 1
    assert pred.shape == label.shape,'the shape of pred and label should be same'
    intersection = pred * (pred == label)
    area_inter, _ = np.histogram(intersection, bins=num_class, range=(1, num_class))
    area_pred, _ = np.histogram(pred, bins=num_class, range=(1, num_class))
    area_label, _ = np.histogram(label, bins=num_class, range=(1, num_class))
    area_union = area_pred + area_label - area_inter
    iou = area_inter / area_union
    return iou[np.newaxis,...]



def mean_iou(input, target, classes = 2):
    """  compute the value of mean iou
    :param input:  2d array, int, prediction
    :param target: 2d array, int, ground truth
    :param classes: int, the number of class
    :return:
        miou: float, the value of miou
    """
    miou = []
    for i in range(classes):
        intersection = np.logical_and(target == i, input == i)
        # print(intersection.any())
        union = np.logical_or(target == i, input == i)
        temp = np.sum(intersection) / np.sum(union)
        miou += [temp]
    return  np.array(miou)[np.newaxis,...]



if __name__ == '__main__':
    all_iou = np.zeros((1,5))
    for index in range(100):
        batch_pred = torch.randint(high=6,low=0,size=(256,256)).numpy()
        batch_label = torch.randint(high=6,low=0,size=(256,256)).numpy()
        miou = mean_iou(batch_pred,batch_label,classes=5)
        print(miou.shape)
        all_iou = np.concatenate((all_iou,miou),axis=0)
        print(all_iou.shape)

